{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据集验证，需要加载的包库\n",
    "import os\n",
    "from PIL import Image\n",
    "from matplotlib import pyplot as plt\n",
    "import xml\n",
    "import cv2 as cv\n",
    "import numpy as np\n",
    "import json\n",
    "import math\n",
    "from scipy.spatial import distance\n",
    "from collections import OrderedDict\n",
    "import time\n",
    "import pandas as pds\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import models\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from utilour.datasetLoader import MyDataset\n",
    "import torch.optim as optimer\n",
    "from logdir import loginfo\n",
    "from utilour.tool import calaccurary,visualize_cam\n",
    "from torchsummary import summary\n",
    "from utilour.gradcadpp import GradCAM\n",
    "import torchvision\n",
    "from torchvision.utils import make_grid, save_image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载数据\n",
    "csvpath=\"/media/gis/data/jupyterlabhub/gitcode/hrx/dataset/train_test_vail.csv\" # 文件路径\n",
    "traindataset=MyDataset(csvpath,\"train\")\n",
    "vaildataset=MyDataset(csvpath,\"vail\")\n",
    "testdataset=MyDataset(csvpath,\"test\")\n",
    "trainLoader=DataLoader(dataset=traindataset,batch_size=8,num_workers=8)\n",
    "vailLoader=DataLoader(dataset=vaildataset,batch_size=8,num_workers=8)\n",
    "testLoader=DataLoader(dataset=testdataset,batch_size=8,num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载模型--修改模型\n",
    "resnet50cpk=torch.load(os.path.join(\".\",\"modelRecord\",\"resnet50.pkl\"))\n",
    "summary(resnet50cpk,input_size=(3,217,217))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 声明gradcam\n",
    "gradcam=GradCAM({\"type\":\"resnet50\",\"model\":resnet50cpk,\"target_layer\":resnet50cpk._modules[\"layer4\"]})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "font=cv.FONT_HERSHEY_SIMPLEX#使用默认字体\n",
    "images=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for step,(data,label,imgpath) in  enumerate(testLoader):\n",
    "    data=data.cuda()\n",
    "    b,c,h,w=data.shape\n",
    "    label=label.cuda()\n",
    "    #tlabels=resnet50cpk(data)\n",
    "    #saliency_maps,tlabel=gradcam.calGradCAM(data)\n",
    "    calres=gradcam.calGradCAM(data)\n",
    "    labellist=label.cpu().numpy()\n",
    "\n",
    "    imglist=data.cpu()\n",
    "    for b in range(len(calres)):\n",
    "        #原始图片\n",
    "        img=imglist[b,:,:,:]\n",
    "\n",
    "        tempimgs=[]\n",
    "        tempinfo=calres[b] # 获取当前的类别计算情况\n",
    "        tl=labellist[b]\n",
    "        reslabel=tempinfo['label'].detach().cpu().numpy()\n",
    "        gl=np.argmax(reslabel)\n",
    "\n",
    "        # 图片的信息\n",
    "        imginfo=np.zeros((h,w,c),np.uint8) # 新建图像，注意一定要是uint8\n",
    "        imginfo=cv.putText(imginfo,'{}'.format(os.path.basename(imgpath[0])),(0,50),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "        imginfo=cv.putText(imginfo,'{}-T-{}'.format(tl,gl),(0,100),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "        imginfo=imginfo/255.0\n",
    "        imginfo=torch.from_numpy(np.transpose(imginfo,[2,0,1]))\n",
    "        tempimgs.append(imginfo)\n",
    "        # 标签信息\n",
    "        labelinfo=np.zeros((h,w,c),np.uint8) # 新建图像，注意一定要是uint8\n",
    "        labelinfo=cv.putText(labelinfo,'{}:{}'.format('0','%.3f'%reslabel[0]),(0,50),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "        labelinfo=cv.putText(labelinfo,'{}:{}'.format('1','%.3f'%reslabel[1]),(0,100),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "        labelinfo=labelinfo/255.0\n",
    "        labelinfo=torch.from_numpy(np.transpose(labelinfo,[2,0,1]))\n",
    "        tempimgs.append(labelinfo)\n",
    "        # 类别渲染情况\n",
    "        for k in range(len(tempinfo[\"saliency\"])):\n",
    "            saliencyinfo=tempinfo[\"saliency\"][k]\n",
    "            \n",
    "            # 热力图信息\n",
    "            sa_info=np.zeros((h,w,c),np.uint8) # 新建图像，注意一定要是uint8\n",
    "            sa_info=cv.putText(sa_info,'{}:{}'.format('min','%.7f'%saliencyinfo[\"min\"]),(0,50),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "            sa_info=cv.putText(sa_info,'{}:{}'.format('max','%.7f'%saliencyinfo[\"max\"]),(0,100),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "            sa_info=cv.putText(sa_info,'{}:{}'.format('lab',k),(0,150),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "            sa_info=cv.putText(sa_info,'{}:{}'.format('conf','%.7f'%reslabel[k]),(0,200),font,1.2,(255,255,255),3)#添加文字，1.2表示字体大小，（0,40）是初始的位置，(255,255,255)表示颜色，2表示粗细\n",
    "            sa_info=sa_info/255.0\n",
    "            sa_info=torch.from_numpy(np.transpose(sa_info,[2,0,1]))\n",
    "            tempimgs.append(sa_info)\n",
    "            \n",
    "            # saliency_map \n",
    "            saliency_map=saliencyinfo[\"data\"]\n",
    "            heatmap, result = visualize_cam(saliency_map.detach().cpu().numpy(), img)\n",
    "            tempimgs.append(img.squeeze().cpu())\n",
    "            tempimgs.append(heatmap)\n",
    "            tempimgs.append(result)\n",
    "        images.append(torch.stack(tempimgs, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = make_grid(torch.cat(images, 0), nrow=10)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "print(data.shape)\n",
    "print(images.shape)\n",
    "imgs=images.numpy()*255\n",
    "imgs=np.transpose(imgs,[1,2,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \n",
    "print( os.listdir( \".\"))\n",
    "output_dir =os.path.join(\".\", 'outputs')\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "output_name = \"test.jpg\"\n",
    "output_path = os.path.join(output_dir, output_name)\n",
    "print(output_path)\n",
    "cv.imwrite(output_path,imgs.astype(np.uint8))\n",
    "d=Image.fromarray(imgs.astype(np.uint8))\n",
    "d.save(os.path.join(output_dir, \"testPIL.jpg\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
